# -*- coding: utf-8 -*-
"""
股票日线回测

1.提供ema每日数据
2.提供每日的 开，高，低，收 四个价格
3.所有数据都是有时间索引的
"""

import pandas as pd
import os

from stock_py.config import config as conf
from stock_py.tools.database import stock_db as mydb

# 定义全局变量
'''
is_buy = 0  # 当前是否购买该股票
last_ema = 0  # 上一收盘价计算出来的ema价格
free = 0  # 当前账户收益
trad_cnt = 0  # 交易次数
'''
# 默认每次交易100股，每次交易费用是5元
'''
global is_buy = 0  # 当前是否购买该股票
global last_ema = 0  # 上一收盘价计算出来的ema价格
global free = 0  # 当前账户收益
global trad_cnt = 0  # 交易次数
'''

file = r"F:\stock_data\overview-data-sz\sz000065.csv"
df = read_data(file)
fee_line = pd.Series(index=range(1,20))
trad_line = pd.Series(index=range(1,20))
for i in range(1,20):
    code, name, trad_fee, fee = calc_inconme(df, n_ema=i)
    fee_line[i] = fee
    trad_line[i] = trad_fee

fee_line.plot()
trad_line.plot()



def read_data(file):
    '''读取csv中的股票数据'''
    data = pd.read_csv(file, encoding="GB18030")
    data = data[['股票代码', '股票名称', '交易日期', '开盘价', '最高价', '最低价', '收盘价', '前复权价']]
    data.columns = ['code', 'name', 'dates', 'open', 'high', 'low', 'close', 'qfq']
    data['dates'] = data['dates'].apply(lambda s: pd.to_datetime(s))
    data = data.sort_values(by='dates')
    data.index = range(len(data))
    return data

def calc_inconme(data, n_ema=15):
    '''根据df提供的基础数据计算ema的收益.
    data 是包含code，date,四个价格的数据库。
    ema表示要计算的均线，会在函数中计算eam的值，然后再计算收益
    '''
    is_buy = 0  # 当前是否购买该股票
    buy_price = 0  # 买入该股票的价格，即成本价
    fee = 0  # 当前账户收益
    trad_fee = 0  # 交易费用
    # ---------------------------------------------------
    # 使用前复权计算ema
    # 注意，ema计算需要xuj
    data['ema'] = pd.ewma(data['close'], span=n_ema)
    # 数据偏移，因为今天用来比较的ema用的是昨天的ema值
    # 所以需要沿着时间轴往后面偏移一天，这样是为了更好的计算，加快速度
    data2 = data.shift(1)
    data['ema'] = data2['ema']
    data = data[n_ema:]
    data.index = range(len(data))
    # ---------------------------------------------------
    # 1.如果已经购买，则看最低价是否低于昨天收盘价计算出来的ema，低于则卖出
    # 2.卖出后如果收盘价大于eam，则要买回来
    # 3.如果没有购买，则看最高家是否大于ema,高于则买入
    # ---------------------------------------------------
    trad_line = pd.Series(index=data.index)
    fee_line = pd.Series(index=data.index)  # 绘制收益曲线需要的数据
    buy_line = pd.Series(index=data.index)

    for i in range(len(data)):
        today_sell = 0
        high = data['high'][i]
        low = data['low'][i]
        close = data['close'][i]
        ema = data['ema'][i]
        # ---------------------------------------------------
        # 如果已经购买，则看最低价是否低于昨天收盘价计算出来的ema，低于则卖出
        if is_buy == 1 and low < ema:
            is_buy = 0    # 卖出
            today_sell = 1  # 今天卖出
            fee = fee + (ema - buy_price) * 100 - 5  # 计算收益
            trad_fee = trad_fee + 1  # 计算交易次数
            fee_line[i] = fee
            trad_line[i] = trad_fee
            buy_line[i] = is_buy
        # ---------------------------------------------------
        # 今天卖出后如果收盘价大于eam，则要买回来
        if today_sell == 1 and close > ema:
            is_buy = 1   # 买入
            buy_price = ema  # 买入价格
            trad_fee = trad_fee + 1
            fee_line[i] = fee
            trad_line[i] = trad_fee
            buy_line[i] = is_buy
            continue
        # ---------------------------------------------------
        # 如果没有购买，则看最高家是否大于ema,高于则买入
        if is_buy == 0 and high > ema:
            is_buy = 1
            trad_fee = trad_fee + 1
            buy_price = ema
            fee_line[i] = fee
            trad_line[i] = trad_fee
            buy_line[i] = is_buy
            continue
        # ---------------------------------------------------
        # 如果今天没有操作，就用昨天的数据填充series
        fee_line[i] = fee
        trad_line[i] = trad_fee
        buy_line[i] = is_buy
        # ---------------------------------------------------
    if is_buy == 1:
        is_buy = 0
        trad_fee = trad_fee + 1
        fee = fee + (ema - buy_price) * 100 - 5
        fee_line[i] = fee
        trad_line[i] = trad_fee
        buy_line[i] = is_buy
    # ---------------------------------------------------
    code = data['code'][i]
    name = data['name'][i]
    return code, name, trad_fee, fee


def calc_all(path):
    '''计算path路径下的所有数据文件的股票收益
    。也就是整合上面的两个函数'''
    # 获取当前路径下的所有数据文件,并转成全路径地址
    files = [file for file in os.listdir(path) if '.csv' in file]
    files = [os.path.join(path, file) for file in files]
    calc_data = []
    file_cnt = len(files)
    for i in range(len(files)):
        try:
            data = read_data(files[i])
            calc = calc_inconme(data, n_ema=4)  # 返回code,name,trad_fee,fee，作为一个元组
            calc_data.append(calc)
        except:
            pass
        print("%d / %d" % (i, file_cnt))
    # 转成df
    df = pd.DataFrame(calc_data, columns=['code', 'name', 'trad_fee', 'fee'])
    return df


def main(path):
    # 导入日志
    logger = conf.get_log()
    logger.info('---------------------------------------')
    logger.info('开始计算股票收益')
    # -------------------------------------------------------
    path1 = r"F:\stock_data\overview-data-sh"
    path2 = r"F:\stock_data\overview-data-sz"
    df1 = calc_all(path1)
    df2 = calc_all(path2)
    df = pd.concat([df1, df2])
    # -------------------------------------------------------
    # 写入excel
    df.to_excel("eam4指标收益回测.xlsx", "所有股票收益回测")

if __name__ == '__main__':
    main()
